import argparse


class ArgumentParser:
    def __init__(self):
        self.parser = argparse.ArgumentParser(description='The PDF book Translate')
        self.parser.add_argument('--config', type=str, default='config.yaml', help='配置文件')
        self.parser.add_argument('--model_type', type=str, required=True, choices=['GLMModel', 'OpenAIModel'],
                                 help='The type of translation model to use. Choose between "GLMModel" and '
                                      '"OpenAIModel".')
        self.parser.add_argument('--glm_model_url', type=str, help='The URL of the ChatGLM model URL.')
        self.parser.add_argument('--timeout', type=int, help='Timeout for the API request in seconds.')
        self.parser.add_argument('--openai_model', type=str,
                                 help='The model name of OpenAI Model. Required if model_type is "OpenAIModel".')
        self.parser.add_argument('--api_key', type=str,
                                 help='The API key for OpenAIModel. Required if model_type is "OpenAIModel".')
        # self.parser.add_argument('--azure_api_endpoint', type=str,
        #                          help='The API key for OpenAIModel. Required if model_type is "api_endpoint".')
        self.parser.add_argument('--book', type=str, help='PDF file to translate.')
        self.parser.add_argument('--file_format', type=str,
                                 help='The file format of translated book. Now supporting PDF and Markdown')

    def arguments_parser(self):
        args = self.parser.parse_args()
        return args

    def check_oepnai_args(self, model_type, api_key, openai_model):
        if model_type == 'OpenAIModel' and not api_key and not openai_model:
            self.parser.error('参数不正确')
            return False
        return True
